/**
 * Copyright (c) 2021-2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ecma_reference_processor.h"
#include "mem/object_helpers.h"
#include "runtime/include/coretypes/class.h"
#include "runtime/mem/gc/gc.h"
#include "plugins/ecmascript/runtime/ecma_vm.h"
#include "plugins/ecmascript/runtime/js_hclass.h"
#include "plugins/ecmascript/runtime/js_weak_container.h"
#include "plugins/ecmascript/runtime/linked_hash_table-inl.h"
#include "plugins/ecmascript/runtime/js_tagged_value-inl.h"

namespace ark::mem::ecmascript {

EcmaReferenceProcessor::EcmaReferenceProcessor(ark::ecmascript::EcmaVM *vm) : vm_(vm), gc_(vm->GetGC()) {}

template <typename Callback>
bool EnumerateArrayElements(const ark::ecmascript::TaggedArray *array, const Callback &cb)
{
    for (size_t i = 0; i < array->GetLength(); ++i) {
        ark::ecmascript::JSTaggedValue key = array->Get(i);
        if (!key.IsHeapObject()) {
            continue;
        }
        if (cb(i, key.GetHeapObject())) {
            return true;
        }
    }
    return false;
}

template <typename Container, typename Callback>
bool EnumerateKeys(Container *container, const Callback &cb)
{
    int totalElements = container->NumberOfElements() + container->NumberOfDeletedElements();
    for (int i = 0; i < totalElements; ++i) {
        ark::ecmascript::JSTaggedValue key = container->GetKey(i);
        if (key.IsHole()) {
            continue;
        }
        ASSERT(key.IsHeapObject());
        if (cb(i, key.GetHeapObject())) {
            return true;
        }
    }
    return false;
}

bool EcmaReferenceProcessor::IsReference([[maybe_unused]] const BaseClass *baseCls, const ObjectHeader *ref,
                                         const ReferenceCheckPredicateT &refPred) const
{
    auto phase = gc_->GetGCPhase();
    ASSERT(IsMarking(phase));
    if (gc_->GetType() == GCType::G1_GC && (phase == GCPhase::GC_PHASE_MARK || phase == GCPhase::GC_PHASE_REMARK) &&
        !gc_->IsFullGC()) {
        // dont process refs on conc-mark in G1, it can cause a high pause
        LOG(DEBUG, REF_PROC) << "Skip reference: " << ref << " because it's G1 with phase: " << static_cast<int>(phase);
        return false;
    }
    ASSERT(ref->ClassAddr<BaseClass>()->IsDynamicClass());

    auto *hcls = ark::ecmascript::JSHClass::FromHClass(ref->ClassAddr<HClass>());
    if (!hcls->IsWeakContainer()) {
        return false;
    }
    auto isReferenceChecker = [this, &refPred](int /* unused */, ObjectHeader *key) {
        return refPred(key) && !gc_->IsMarked(key);
    };
    auto objectType = hcls->GetObjectType();
    switch (objectType) {
        case ark::ecmascript::JSType::TAGGED_ARRAY:
            // Weak tagged array is used for inline caches.
            // It is reasonable to collect IC only in OOM case else it leads
            // to big perf degradation.
            if (gc_->GetLastGCCause() == GCTaskCause::OOM_CAUSE) {
                return EnumerateArrayElements(static_cast<const ark::ecmascript::TaggedArray *>(ref),
                                              isReferenceChecker);
            }
            break;
        case ark::ecmascript::JSType::JS_WEAK_REF: {
            ark::ecmascript::JSTaggedValue referent =
                static_cast<const ark::ecmascript::JSWeakRef *>(ref)->GetReferent();
            return referent.IsHeapObject() ? isReferenceChecker(0, referent.GetRawHeapObject()) : false;
        }
        case ark::ecmascript::JSType::LINKED_HASH_MAP:
            return EnumerateKeys(static_cast<const ark::ecmascript::LinkedHashMap *>(ref), isReferenceChecker);
        case ark::ecmascript::JSType::LINKED_HASH_SET:
            return EnumerateKeys(static_cast<const ark::ecmascript::LinkedHashSet *>(ref), isReferenceChecker);
        default:
            LOG(FATAL, REF_PROC) << "Unknown weak container";
    }
    return false;
}

void EcmaReferenceProcessor::HandleReference([[maybe_unused]] GC *gc, [[maybe_unused]] GCMarkingStackType *objectsStack,
                                             [[maybe_unused]] const BaseClass *baseClass, const ObjectHeader *object,
                                             [[maybe_unused]] const ReferenceProcessPredicateT &pred)
{
    os::memory::LockHolder lock(weakRefLock_);
    dynWeakReferences_.insert(const_cast<ObjectHeader *>(object));
}

void EcmaReferenceProcessor::ProcessReferences([[maybe_unused]] bool concurrent,
                                               [[maybe_unused]] bool clearSoftReferences,
                                               [[maybe_unused]] GCPhase gcPhase,
                                               const mem::GC::ReferenceClearPredicateT &pred)
{
    os::memory::LockHolder lock(weakRefLock_);
    ark::ecmascript::JSThread *thread = vm_->GetAssociatedJSThread();
    while (!dynWeakReferences_.empty()) {
        ObjectHeader *reference = *dynWeakReferences_.begin();
        dynWeakReferences_.erase(*dynWeakReferences_.begin());
        auto *hcls = ark::ecmascript::JSHClass::FromHClass(reference->ClassAddr<HClass>());
        ASSERT(hcls->IsWeakContainer());
        auto objectType = hcls->GetObjectType();
        if (objectType == ark::ecmascript::JSType::TAGGED_ARRAY) {
            auto *array = static_cast<ark::ecmascript::TaggedArray *>(reference);
            auto handler = [this, array](uint32_t index, ObjectHeader *elem) {
                if (!gc_->IsMarked(elem)) {
                    ObjectAccessor::SetDynValueWithoutBarrier(array, array->GetElementOffset(index),
                                                              ark::ecmascript::JSTaggedValue::Undefined().GetRawData());
                }
                return false;
            };
            EnumerateArrayElements(array, handler);
        } else if (objectType == ark::ecmascript::JSType::JS_WEAK_REF) {
            auto *ref = static_cast<ark::ecmascript::JSWeakRef *>(reference);
            ASSERT(ref->GetReferent().IsHeapObject());
            if (!gc_->IsMarked(ref->GetReferent().GetRawHeapObject())) {
                ObjectAccessor::SetDynValueWithoutBarrier(ref, ark::ecmascript::JSWeakRef::GetReferentOffset(),
                                                          ark::ecmascript::JSTaggedValue::Undefined().GetRawData());
            }
        } else if (objectType == ark::ecmascript::JSType::LINKED_HASH_MAP) {
            auto *map = static_cast<ark::ecmascript::LinkedHashMap *>(reference);
            auto mapHandler = [this, map, thread](int index, ObjectHeader *key) {
                if (!gc_->IsMarked(key)) {
                    map->RemoveEntry(thread, index);
                }
                return false;
            };
            EnumerateKeys(map, mapHandler);
        } else if (objectType == ark::ecmascript::JSType::LINKED_HASH_SET) {
            auto *set = static_cast<ark::ecmascript::LinkedHashSet *>(reference);
            auto setHandler = [this, set, thread](int index, ObjectHeader *key) {
                if (!gc_->IsMarked(key)) {
                    set->RemoveEntry(thread, index);
                }
                return false;
            };
            EnumerateKeys(set, setHandler);
        } else {
            LOG(FATAL, REF_PROC) << "Unknown weak container";
        }
    }

    ark::ecmascript::WeakRootVisitor gcUpdateWeak =
        [this, pred](ark::ecmascript::TaggedObject *header) -> ark::ecmascript::TaggedObject * {
        if (pred(header) && !gc_->IsMarked(header)) {
            return nullptr;
        }

        return header;
    };
    vm_->ProcessReferences(gcUpdateWeak);
}

}  // namespace ark::mem::ecmascript
